package cn.code.controller;

import cn.code.pojo.ModelReference;
import cn.code.pojo.ResponseParser;
import cn.code.service.FileProcessingService;
import cn.code.util.ToolsUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@RestController
@RequestMapping("/api/chat")
@CrossOrigin
@Slf4j
public class ChatController {


    @Autowired
    private ThreadPoolTaskExecutor  threadPoolTaskExecutor;

    private final OpenAiChatModel chatModel;
    private final FileProcessingService fileProcessingService;
    //private final VectorDBService vectorDBService;
    private final ExecutorService executor = Executors.newCachedThreadPool();

    private final ConcurrentHashMap<String, SseEmitter> emitters = new ConcurrentHashMap<>();

    @Value("${app.file.upload-dir}")
    private String uploadDir;

    @Autowired
    public ChatController(OpenAiChatModel chatModel,
                          FileProcessingService fileProcessingService
                          //VectorDBService vectorDBService
    ) {
        this.chatModel = chatModel;
        this.fileProcessingService = fileProcessingService;
        //this.vectorDBService = vectorDBService;
    }

    @GetMapping("/ai/generate")
    public Map generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return Map.of("generation", this.chatModel.call(message));
    }
    private List<ChatResponse> createFollowUpQuestions() {
        String message = "您是否还需要其他帮助？";
        List<Generation> generations = Arrays.asList(
            new Generation(new AssistantMessage(message)),
            new Generation(new AssistantMessage(message)),
            new Generation(new AssistantMessage(message))
        );
        return List.of(new ChatResponse(generations));
    }

    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {

        // 准备后续问题
        Flux<ChatResponse> followUpQuestions = Flux.fromIterable(createFollowUpQuestions());

        //Prompt prompt1 = new Prompt(
        //        "Generate the names of 5 famous pirates.",
        //        OpenAiChatOptions.builder()
        //                .model("deepseek-chat")
        //                .temperature(0.4)
        //                .build()
        //);
        // 输出Flux里面的属性，我需要判断当前Flux是否快要结束了
        Prompt prompt = new Prompt(new UserMessage(message));
        Flux<ChatResponse> stream = this.chatModel.stream(prompt);

        stream = stream.concatWith(followUpQuestions);

        // 使用doOnNext来记录日志，这样不会影响数据流
        stream = stream.doOnNext(chunk -> log.debug("chunk：{}", chunk != null ? chunk : "null chunk"));

        return stream;
    }




    // SSE实现方式
    @GetMapping(path = "/ai/generateStream/log", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter subscribe2( @RequestParam(value = "message",
                                          defaultValue = "Tell me a joke") String message
            , @RequestParam(value = "model", required = false) String model) {

        // 超时时间小于30秒
        SseEmitter emitter = new SseEmitter(30_000L);
        executor.execute(() -> {
            ToolCallback[] tools = ToolCallbacks.from(new ToolsUtil());
            ChatOptions chatOptions = ToolCallingChatOptions.builder()
                    .toolCallbacks(tools)
                    .build();
            Prompt prompt = new Prompt(message, chatOptions);

            //打印引用
            Flux<ChatResponse> stream = this.chatModel.stream(prompt)
                    .doOnNext(chunk -> {
                        log.debug("chunk: " + chunk);
                        // 实时解析引用
                        List<ModelReference> refs = ResponseParser.extractReferences(chunk.toString());
                        if (!refs.isEmpty()) {
                            refs.forEach(ref ->
                                    log.debug("发现引用: {} - {}", ref.getIndex(), ref.getTitle())
                            );
                        }
                    });



            // 打印数据流
            stream.subscribe(chunk -> {
                try {

                    emitter.send(chunk);
                } catch (Exception e) {
                    emitter.completeWithError(e);
                }
            });
        });
        return emitter;
    }



    @GetMapping(path = "/ai/generateStream/data", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter subscribe1( @RequestParam(value = "message",
                                          defaultValue = "Tell me a joke") String message
            , @RequestParam(value = "model", required = false) String model) {

        // 超时时间小于30秒
        SseEmitter emitter = new SseEmitter(30_000L);
        executor.execute(() -> {
            ToolCallback[] tools = ToolCallbacks.from(new ToolsUtil());
            ChatOptions chatOptions = ToolCallingChatOptions.builder()
                    .toolCallbacks(tools)
                    .build();
            Prompt prompt = new Prompt(message, chatOptions);
            Flux<String> stream = this.chatModel.stream(prompt.toString());


            // 打印数据流
            stream.subscribe(chunk -> {
                try {

                    emitter.send(chunk);
                } catch (Exception e) {
                    emitter.completeWithError(e);
                }
            });
        });
        return emitter;
    }


    /**
     * 文件上传接口
     */
    @PostMapping("/upload")
    public ResponseEntity<Map<String, Object>> uploadFile(@RequestParam("filespm") MultipartFile[] files) {
        Map<String, Object> response = new HashMap<>();
        MultipartFile file = files[0];
        try {
            // 检查文件类型
            if (!fileProcessingService.isSupportedFileType(file.getOriginalFilename())) {
                response.put("success", false);
                response.put("message", "不支持的文件类型");
                return ResponseEntity.badRequest().body(response);
            }

            // 保存文件
            String fileName = System.currentTimeMillis() + "_" + file.getOriginalFilename();
            Path uploadPath = Paths.get(uploadDir);
            if (!Files.exists(uploadPath)) {
                Files.createDirectories(uploadPath);
            }

            Path filePath = uploadPath.resolve(fileName);
            Files.copy(file.getInputStream(), filePath);

            // 提取文件内容
            String content = fileProcessingService.extractContent(file);
            String fileType = fileProcessingService.getFileType(file.getOriginalFilename());


            // TODO: 生成向量嵌入并存储到向量数据库
            // List<Float> embeddings = generateEmbeddings(content);
            // vectorDBService.insertDocument(content, embeddings, fileType, fileName);

            response.put("success", true);
            response.put("message", "文件上传成功");
            response.put("fileName", fileName);
            response.put("fileType", fileType);
            response.put("filePath", filePath.toString());
            response.put("contentLength", content.length());

            return ResponseEntity.ok(response);

        } catch (Exception e) {
            response.put("success", false);
            response.put("message", "文件上传失败: " + e.getMessage());
            return ResponseEntity.internalServerError().body(response);
        }
    }

    /**
     * 搜索文档接口
     */
    @GetMapping("/search")
    public ResponseEntity<Map<String, Object>> searchDocuments(@RequestParam("query") String query,
                                                               @RequestParam(value = "topK", defaultValue = "5") int topK) {
        Map<String, Object> response = new HashMap<>();

        try {
            // TODO: 生成查询向量的嵌入
            // List<Float> queryEmbedding = generateEmbeddings(query);

            // 搜索相似文档
            // List<VectorDBService.SearchResult> results = vectorDBService.searchSimilarDocuments(queryEmbedding, topK);

            List<Map<String, Object>> mockResults = new ArrayList<>();
            // 模拟搜索结果
            for (int i = 0; i < Math.min(topK, 3); i++) {
                Map<String, Object> result = new HashMap<>();
                result.put("content", "这是第 " + (i + 1) + " 个搜索结果，查询词: " + query);
                result.put("fileType", "文档");
                result.put("fileName", "example_" + (i + 1) + ".txt");
                result.put("score", 0.9 - (i * 0.1));
                mockResults.add(result);
            }

            response.put("success", true);
            response.put("results", mockResults);
            response.put("query", query);
            response.put("total", mockResults.size());

            return ResponseEntity.ok(response);

        } catch (Exception e) {
            response.put("success", false);
            response.put("message", "搜索失败: " + e.getMessage());
            return ResponseEntity.internalServerError().body(response);
        }
    }

    /**
     * MCP工具调用接口
     */
    @PostMapping("/mcp/tool")
    public ResponseEntity<Map<String, Object>> callMcpTool(@RequestBody Map<String, Object> request) {
        Map<String, Object> response = new HashMap<>();
        String toolName = (String) request.get("tool");
        Map<String, Object> arguments = (Map<String, Object>) request.get("arguments");

        try {
            switch (toolName) {
                case "search_documents":
                    String query = (String) arguments.get("query");
                    int topK = arguments.get("topK") != null ? (Integer) arguments.get("topK") : 5;
                    // 调用搜索功能
                    return searchDocuments(query, topK);

                case "upload_file":
                    // 这里需要处理文件上传，但需要通过multipart/form-data
                    response.put("success", false);
                    response.put("message", "请使用/upload接口进行文件上传");
                    return ResponseEntity.badRequest().body(response);

                default:
                    response.put("success", false);
                    response.put("message", "未知的工具: " + toolName);
                    return ResponseEntity.badRequest().body(response);
            }
        } catch (Exception e) {
            response.put("success", false);
            response.put("message", "工具调用失败: " + e.getMessage());
            return ResponseEntity.internalServerError().body(response);
        }
    }

    /**
     * 获取支持的模型列表
     */
    @GetMapping("/models")
    public ResponseEntity<Map<String, Object>> getModels() {
        Map<String, Object> response = new HashMap<>();
        List<Map<String, String>> models = new ArrayList<>();

        models.add(Map.of("id", "gpt-3.5-turbo", "name", "GPT-3.5 Turbo"));
        models.add(Map.of("id", "gpt-4", "name", "GPT-4"));
        models.add(Map.of("id", "deepseek-chat", "name", "DeepSeek Chat"));
        models.add(Map.of("id", "claude-2", "name", "Claude 2"));

        response.put("success", true);
        response.put("models", models);

        return ResponseEntity.ok(response);
    }

    /**
     * 获取上传的文件列表
     */
    @GetMapping("/files")
    public ResponseEntity<Map<String, Object>> getFiles() {
        Map<String, Object> response = new HashMap<>();
        List<Map<String, String>> files = new ArrayList<>();

        try {
            Path uploadPath = Paths.get(uploadDir);
            if (Files.exists(uploadPath)) {
                Files.list(uploadPath)
                        .filter(Files::isRegularFile)
                        .forEach(file -> {
                            Map<String, String> fileInfo = new HashMap<>();
                            fileInfo.put("name", file.getFileName().toString());
                            fileInfo.put("size", String.valueOf(file.toFile().length()));
                            fileInfo.put("type", fileProcessingService.getFileType(file.getFileName().toString()));
                            files.add(fileInfo);
                        });
            }

            response.put("success", true);
            response.put("files", files);

        } catch (IOException e) {
            response.put("success", false);
            response.put("message", "获取文件列表失败: " + e.getMessage());
            return ResponseEntity.internalServerError().body(response);
        }

        return ResponseEntity.ok(response);
    }

    // TODO: 实现向量嵌入生成方法
    private List<Float> generateEmbeddings(String text) {
        // 这里需要调用嵌入模型生成向量
        // 暂时返回模拟数据
        List<Float> embeddings = new ArrayList<>();
        for (int i = 0; i < 1536; i++) {
            embeddings.add((float) Math.random());
        }
        return embeddings;
    }
}
